import {
  WebSocketGateway,
  WebSocketServer,
  SubscribeMessage,
  OnGatewayConnection,
  OnGatewayDisconnect
} from '@nestjs/websockets';
import { Server, Socket } from 'socket.io';
import { Logger } from '@nestjs/common';
import { ComfyUIService } from '../comfyui/comfyui.service';

@WebSocketGateway({
  namespace: 'generation',
  cors: {
    origin: true,
    credentials: true
  }
})
export class GenerationGateway implements OnGatewayConnection, OnGatewayDisconnect {
  private readonly logger = new Logger(GenerationGateway.name);
  private taskSubscriptions = new Map<string, Set<string>>();

  @WebSocketServer()
  server: Server;

  constructor(private readonly comfyuiService: ComfyUIService) { }

  handleConnection(client: Socket) {
    this.logger.log(`Client connected: ${client.id}`);
  }

  handleDisconnect(client: Socket) {
    this.logger.log(`Client disconnected: ${client.id}`);
    // 清理该客户端的所有订阅
    this.cleanupClientSubscriptions(client.id);
  }

  @SubscribeMessage('subscribe_task')
  handleSubscribeTask(client: Socket, promptId: string) {
    this.logger.log(`Client ${client.id} subscribed to task ${promptId}`);

    // 记录订阅关系
    if (!this.taskSubscriptions.has(promptId)) {
      this.taskSubscriptions.set(promptId, new Set());
    }
    this.taskSubscriptions.get(promptId)?.add(client.id);

    // 监听 ComfyUI 的状态更新
    this.comfyuiService.onTaskUpdate(promptId, (data) => {
      this.handleTaskUpdate(promptId, data);
    });
  }

  private handleTaskUpdate(promptId: string, data: any) {
    //console.log(`Task update for ${promptId}:`, data);
    const subscribers = this.taskSubscriptions.get(promptId);
    if (!subscribers) return;

    // 根据状态类型发送不同的事件
    switch (data.type) {
      case 'execution_start':
      case 'progress':
      case 'executing':
      case 'execution_cached':
        this.server.to(Array.from(subscribers)).emit('generation_progress', {
          promptId,
          progress: data.progress,
          comfyui_data: data.data?.comfyui_data || data.comfyui_data // 确保 comfyui_data 正确传递
        });
        break;

      case 'executed':
      case 'execution_success':
        console.log("======completed ===========data:" + data);
        this.server.to(Array.from(subscribers)).emit('generation_complete', {
          promptId,
          outputs: data.outputs,
          comfyui_data: data.comfyui_data // 传递 comfyui_data
        });
        // 清理订阅
        this.cleanupTaskSubscriptions(promptId);
        break;

      case 'error':
        this.server.to(Array.from(subscribers)).emit('generation_error', {
          promptId,
          message: data.error,
          comfyui_data: data.comfyui_data // 传递 comfyui_data
        });
        // 清理订阅
        this.cleanupTaskSubscriptions(promptId);
        break;
    }
  }

  private cleanupClientSubscriptions(clientId: string) {
    for (const [promptId, subscribers] of this.taskSubscriptions.entries()) {
      subscribers.delete(clientId);
      if (subscribers.size === 0) {
        this.taskSubscriptions.delete(promptId);
      }
    }
  }

  private cleanupTaskSubscriptions(promptId: string) {
    this.taskSubscriptions.delete(promptId);
  }
} 